# Copyright (c) Adithya Kumar, The Pennsylvania State University. All rights reserved.
# Licensed under the MIT License.

PROJECT('nnfusion_lstm' LANGUAGES CXX C CUDA)
CMAKE_MINIMUM_REQUIRED(VERSION 3.10)
SET(CMAKE_CXX_STANDARD 11)
SET(CMAKE_CXX_STANDARD_REQUIRED ON)
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
  set(CMAKE_CUDA_ARCHITECTURES $ENV{CUDA_ARCH_CODE_GEN})
endif()
SET(CMAKE_POSITION_INDEPENDENT_CODE ON)
SET(CUDA_SEPARABLE_COMPILATION ON)

INCLUDE_DIRECTORIES(../../utils/ ../../)

SET(NNFUSION_SRC "nnfusion_rt.cu" CACHE STRING "codegen source file")
SET(NNFUSION_TARGET_NAME "nnfusion_lstm_rt" CACHE STRING "codegen target name")

FIND_PACKAGE(CUDA 10 REQUIRED)
FIND_LIBRARY(CUDA_DRIVER_LIBRARY
             NAMES cuda_driver cuda
             HINTS ${CUDA_TOOLKIT_ROOT_DIR}
                   ENV CUDA_PATH
             PATH_SUFFIXES nvidia/current lib64 lib/x64 lib lib64/stubs lib/x64/stubs lib/stubs stubs compat)

find_path(CUDNN_INCLUDE_DIR cudnn.h
    HINTS ${CUDA_TOOLKIT_ROOT_DIR}
    PATH_SUFFIXES cuda/include include)

find_library(CUDNN_LIBRARY cudnn
    HINTS ${CUDA_TOOLKIT_ROOT_DIR}
    PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64)

ADD_LIBRARY(${NNFUSION_TARGET_NAME} STATIC ${NNFUSION_SRC})

message(STATUS "CUDA_INCLUDE_DIRS: ${CUDA_INCLUDE_DIRS}")
message(STATUS "CUDNN_INCLUDE_DIRS: ${CUDNN_INCLUDE_DIR}")

TARGET_INCLUDE_DIRECTORIES(${NNFUSION_TARGET_NAME} PUBLIC ${CUDA_INCLUDE_DIRS} ${CUDNN_INCLUDE_DIR})
#TARGET_LINK_LIBRARIES(${NNFUSION_TARGET_NAME} ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} ${CUDNN_LIBRARIES})

SET(GDR_INCLUDE_DIRS "$ENV{GDR_PATH}/include/")
TARGET_INCLUDE_DIRECTORIES(${NNFUSION_TARGET_NAME} PUBLIC ${GDR_INCLUDE_DIRS})
SET_TARGET_PROPERTIES(${NNFUSION_TARGET_NAME} PROPERTIES CUDA_RESOLVE_DEVICE_SYMBOLS ON)
SET_TARGET_PROPERTIES(${NNFUSION_TARGET_NAME} PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
